# Author   : zhen.fang
# Time     : 2021/1/6 00:57
# Email    : fangeddie@163.com
from datetime import datetime, timedelta

import jwt


class JWTBox:
    def __init__(self, JWT_SECRET_KEY, JWT_EXPIRE_DATE):
        self.JWT_SECRET_KEY = JWT_SECRET_KEY
        self.JWT_EXPIRE_DATE = self.formatExpireDate(JWT_EXPIRE_DATE)

    def jwtEncode(self, data):
        """生成jwt-token"""
        payload = {
            'exp': datetime.utcnow() + timedelta(seconds=self.JWT_EXPIRE_DATE),  # 单位秒
            'iat': datetime.utcnow(),
            'data': data if isinstance(data, dict) else {}
        }
        encoded_jwt = jwt.encode(payload, self.JWT_SECRET_KEY, algorithm='HS256')
        return encoded_jwt

    def jwtDecode(self, encoded_jwt):
        de_code = jwt.decode(encoded_jwt, self.JWT_SECRET_KEY, algorithms=['HS256'])
        return de_code

    @staticmethod
    def formatExpireDate(JWT_EXPIRE_DATE_STR):
        suffix = JWT_EXPIRE_DATE_STR[-1]
        if suffix == 's':
            return int(JWT_EXPIRE_DATE_STR[:-1])
        elif suffix == 'm':
            return int(JWT_EXPIRE_DATE_STR[:-1]) * 60
        elif suffix == 'h':
            return int(JWT_EXPIRE_DATE_STR[:-1]) * 60 * 60
        elif suffix == 'd':
            return int(JWT_EXPIRE_DATE_STR[:-1]) * 60 * 60 * 24
        elif suffix == 'w':
            return int(JWT_EXPIRE_DATE_STR[:-1]) * 60 * 60 * 24 * 7
        else:
            return 60


if __name__ == '__main__':
    jw = JWTBox("question", '2')
    data = {"pw": 1222}
    j = jw.jwtEncode(data)
    print(j)
    print(jw.jwtDecode(j))
